from data_selection import CachedHashedNgramDSIR

raw_datasets = ['/share/projset/multimodel-dsir/llavasft_stats.jsonl']
target_datasets = ['/share/projset/multimodel-dsir/llavasft_stats.jsonl']
from typing import List, Optional, Dict, Callable, Union, Iterable

def parse_example_fn(ex: Dict) -> str:
    """Default parse function from example dict to string

    Args:
        ex (Dict): example dict
    """
    return ' '.join(sortd(ex['labels'].keys()))





dsir = CachedHashedNgramDSIR(raw_datasets, target_datasets, target_parse_example_fn=parse_example_fn, ngrams=1, num_buckets=4000, num_proc=5)
dsir.fit_importance_estimator(num_tokens_to_fit='all')
print(dsir.log_diff)
dsir.cached_compute_importance_weights()
dsir.resample(out_dir=dsir.cache_dir / 'resampled', num_to_sample=100)